#include <immintrin.h>
#include <math.h>
#include <stdio.h>
#include <time.h>
#include "utils.h"

void scalar_nested_conditions(const float *a, const float *b, const float *c,
                              const float *d, const float *e, const float *f,
                              const float *g, const float *h, float *dst, int n)
{

    for (int i = 0; i < n; i++)
    {
        float sa = a[i];
        float sb = b[i];
        float sc = c[i];
        float sd = d[i];
        float se = e[i];
        float sf = f[i];
        float sg = g[i];
        float sh = h[i];
        float final = 0;
        if (sa > 4)
        { // 第一层
            if (sb < 5)
            { // 第二层
                if (sc == 10)
                { // 第三层
                    if (sd != 0)
                    { // 第四层
                        final = sa + sb;
                    }
                    else if (se < 7)
                    { // 第四层else if
                        final = sa - se;
                    }
                    else
                    { // 第四层else
                        final = sa - sc;
                    }
                }
                else if (sf >= -2.0f)
                { // 第三层else if
                    final = sa * sf;
                }
                else
                { // 第三层else
                    final = sa * sb;
                }
            }
            else if (sg > 0)
            { // 第二层else if
                final = sa / sg;
            }
            else
            { // 第二层else
                final = sa / sb;
            }
        }
        else if (sh != 0)
        { // 第一层else if
            final = sqrt(sa + sh);
        }
        else
        { // 第一层else
            final = sqrt(sa);
        }
        dst[i] = final;
    }
}

void simd_nested_conditions(const float *a, const float *b, const float *c,
                            const float *d, const float *e, const float *f,
                            const float *g, const float *h, float *dst, int n)
{
    const __m256 zero = _mm256_setzero_ps();
    const __m256 seven = _mm256_set1_ps(7.0);
    const __m256 ten = _mm256_set1_ps(10.0);
    const __m256 neg_two = _mm256_set1_ps(-2.0);
    const __m256 five = _mm256_set1_ps(5.0);
    const __m256 four = _mm256_set1_ps(4.0);
    for (int i = 0; i < n; i += 8)
    {
        // 加载所有输入数据
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        __m256 vc = _mm256_loadu_ps(c + i);
        __m256 vd = _mm256_loadu_ps(d + i);
        __m256 ve = _mm256_loadu_ps(e + i);
        __m256 vf = _mm256_loadu_ps(f + i);
        __m256 vg = _mm256_loadu_ps(g + i);
        __m256 vh = _mm256_loadu_ps(h + i);

        // 生成条件掩码（从最内层到外层）
        // 第四层条件
        __m256 mask_d_ne_zero = _mm256_cmp_ps(vd, zero, _CMP_NEQ_OQ); // d != 0
        __m256 mask_e_lt_7 = _mm256_cmp_ps(ve, seven, _CMP_LT_OQ);    // e < 7

        // 第三层条件
        __m256 mask_c_eq_10 = _mm256_cmp_ps(vc, ten, _CMP_EQ_OQ);     // c == 10
        __m256 mask_f_ge_n2 = _mm256_cmp_ps(vf, neg_two, _CMP_GE_OQ); // f >= -2

        // 第二层条件
        __m256 mask_b_lt_5 = _mm256_cmp_ps(vb, five, _CMP_LT_OQ); // b < 5
        __m256 mask_g_gt_0 = _mm256_cmp_ps(vg, zero, _CMP_GT_OQ); // g > 0

        // 第一层条件
        __m256 mask_a_gt_4 = _mm256_cmp_ps(va, four, _CMP_GT_OQ);  // a > 4
        __m256 mask_a_ne_0 = _mm256_cmp_ps(vh, zero, _CMP_NEQ_OQ); // h != 0

        // 预计算所有可能的分支结果
        __m256 add_ab = _mm256_add_ps(va, vb);
        __m256 sub_ae = _mm256_sub_ps(va, ve);
        __m256 sub_ac = _mm256_sub_ps(va, vc);
        __m256 mul_af = _mm256_mul_ps(va, vf);
        __m256 mul_ab = _mm256_mul_ps(va, vb);
        __m256 div_ag = _mm256_div_ps(va, vg);
        __m256 div_ab = _mm256_div_ps(va, vb);
        __m256 sqrt_ah = _mm256_sqrt_ps(_mm256_add_ps(va, vh));
        __m256 sqrt_a = _mm256_sqrt_ps(va);

        // 从最内层开始逐级混合
        // 第四层混合
        __m256 cond3 = _mm256_and_ps(mask_b_lt_5, mask_a_gt_4);
        __m256 cond4 = _mm256_and_ps(mask_c_eq_10, cond3);
        __m256 layer4_elseif = _mm256_blendv_ps(sub_ac, sub_ae, mask_e_lt_7); // e < 2 ? sub_ae : sub_ac
        __m256 layer4 = _mm256_blendv_ps(layer4_elseif, add_ab, mask_d_ne_zero);

        // 第三层混合
        __m256 layer3_elseif = _mm256_blendv_ps(mul_ab, mul_af, mask_f_ge_n2); // f >=1 ? mul_af : mul_ab
        __m256 layer3 = _mm256_blendv_ps(layer3_elseif, layer4, mask_c_eq_10);

        // 第二层混合
        __m256 cond2 = mask_a_gt_4;
        __m256 layer2_elseif = _mm256_blendv_ps(div_ab, div_ag, mask_g_gt_0); // g !=0 ? div_ag : div_ab
        __m256 layer2 = _mm256_blendv_ps(layer2_elseif, layer3, mask_b_lt_5);

        // 第一层混合
        __m256 final_elseif = _mm256_blendv_ps(sqrt_a, sqrt_ah, mask_a_ne_0); // h >0 ? sqrt_ah : sqrt_a
        __m256 final = _mm256_blendv_ps(final_elseif, layer2, mask_a_gt_4);

        _mm256_storeu_ps(dst + i, final);
    }
}

int main()
{
    int len = 1;
    for (len = 1000; len < 100000000; len *= 10)
    {
        float *a = aligned_alloc(32, len * sizeof(float));
        float *b = aligned_alloc(32, len * sizeof(float));
        float *c = aligned_alloc(32, len * sizeof(float));
        float *d = aligned_alloc(32, len * sizeof(float));
        float *e = aligned_alloc(32, len * sizeof(float));
        float *f = aligned_alloc(32, len * sizeof(float));
        float *g = aligned_alloc(32, len * sizeof(float));
        float *h = aligned_alloc(32, len * sizeof(float));
        float *dst_scalar = aligned_alloc(32, len * sizeof(float));
        float *dst_simd = aligned_alloc(32, len * sizeof(float));

        // 初始化数据
        for (int i = 0; i < len; ++i)
        {
            a[i] = (float)rand() / RAND_MAX;
            b[i] = (float)rand() / RAND_MAX;
            c[i] = (float)rand() / RAND_MAX;
            d[i] = (float)rand() / RAND_MAX;
            e[i] = (float)rand() / RAND_MAX;
            f[i] = (float)rand() / RAND_MAX;
            g[i] = (float)rand() / RAND_MAX;
            h[i] = (float)rand() / RAND_MAX;
        }
        printf("数据个数 %d\n", len);
        // 标量版本计时
        BENCHMARK_FUNC(10, 100, scalar_nested_conditions, a, b, c, d, e, f, g, h, dst_scalar, len);

        // SIMD版本计时
        BENCHMARK_FUNC(10, 100, simd_nested_conditions, a, b, c, d, e, f, g, h, dst_simd, len);

        // 校验结果一致性
        for (int i = 0; i < len; ++i)
        {
            if (dst_scalar[i] != dst_simd[i])
            {
                printf("dst mismatch at index %d\n", i);
                break;
            }
        }

        free(a);
        free(b);
        free(c);
        free(d);
        free(e);
        free(f);
        free(g);
        free(h);
        free(dst_scalar);
        free(dst_simd);
    }
    return 0;
}